/**************************************************************************************
*
*	Copyright (C) 2010 Brent Andrew O. Chua - See LICENSE.txt for more info.
*		
*	Filename	-	Shader.cpp
*
*	Comments	-	See Shader.h
*
**************************************************************************************/

#ifdef __PSX_WINDOWS_WINDOW

#include "../Include/Shader.h"
#include "../Include/Renderer.h"
#include "../Include/Engine.h"
#include "../Include/FileSystem.h"
#include "../Include/ResourceCache.h"
#include "../Include/DX9Mapping.h"
#include "../Include/ShaderConstant.h"

namespace Pulse
{
	Shader::Shader( void )
		: m_pVertexShader( PSX_NULL ), m_pConstants( PSX_NULL )
	{
		RendererDevice *pDevice = Renderer::GetInstance()->GetDevice();
		m_pDevice = reinterpret_cast<IDirect3DDevice9*>(Renderer::GetInstance()->GetNakedDevice());

	}

	//Shader::Shader( ShaderDescription *pDesc )
	//	: m_shaderDesc( *pDesc ), m_pVertexShader( PSX_NULL ), m_pConstants( PSX_NULL )
	//{

	//}

	Shader::~Shader( void )
	{
		DestroyResource();
	}

	BOOL Shader::CreateResource( void )
	{
		// TODO: Finish this function
		return TRUE;
	}

	BOOL Shader::DestroyResource( void )
	{
		// TODO: Finish this function
		return TRUE;
	}

	BOOL Shader::DisableResource( void )
	{
		switch ( m_shaderDesc.type )
		{
		case EShaderType::VERTEX:
			//m_pVertexShader->
			break;
		case EShaderType::GEOMETRY:
			break;
		case EShaderType::PIXEL:
			break;
		case EShaderType::EFFECT:
			m_pEffect->OnLostDevice();
			break;
		}

		return TRUE;
	}

	BOOL Shader::RestoreResource( void )
	{
		switch ( m_shaderDesc.type )
		{
		case EShaderType::VERTEX:
			//m_pVertexShader->
			break;
		case EShaderType::GEOMETRY:
			break;
		case EShaderType::PIXEL:
			break;
		case EShaderType::EFFECT:
			m_pEffect->OnResetDevice();
			break;
		}

		return TRUE;
	}

	EErrorCode::Type Shader::Create( ShaderDescription *pDesc )
	{
		m_shaderDesc = *pDesc;

		switch ( m_shaderDesc.type )
		{
		case EShaderType::VERTEX:
			return CreateVertexShader();
		case EShaderType::GEOMETRY:
			return CreateGeometryShader();
		case EShaderType::PIXEL:
			return CreatePixelShader();
		case EShaderType::EFFECT:
			return CreateEffect();
		default:
			PSX_PushError( "Undefined Shader description type." );
		}

		return EErrorCode::UNKNOWN;
	}

	void Shader::Update( void )
	{
		// Update variables
		SIZE_T numVars = m_variables.GetSize();
		ShaderConstant *pConstant;

		if ( numVars == 0 )
			return;

		for ( SIZE_T i = 0; i < numVars; ++i )
		{
			pConstant = m_variables[i];

			if ( pConstant->IsDynamic() )
			{
				pConstant->Update();

				/*switch( pConstant->GetType() )
				{
				case EShaderConstantType::FLOAT:
					SetFloatArray( pConstant->GetName(), 
						(float*)pConstant->GetDefaultData(), 
						pConstant->GetDataCount() );
					break;
				case EShaderConstantType::INT:
					SetIntegerArray( pConstant->GetName(), 
						(int*)pConstant->GetDefaultData(),
						pConstant->GetDataCount() );
					break;
				case EShaderConstantType::BOOL:
					SetBooleanArray( pConstant->GetName(), 
						(bool*)pConstant->GetDefaultData(), 
						pConstant->GetDataCount() );
					break;
				}*/
			}
		}
	}

	void Shader::Bind( void ) const
	{
		SetDefaults();

		switch( m_shaderDesc.type )
		{
		case EShaderType::VERTEX:
			m_pDevice->SetVertexShader( m_pVertexShader );
			break;
		case EShaderType::PIXEL:
			m_pDevice->SetPixelShader( m_pPixelShader );
			break;
		case EShaderType::EFFECT:
			break;
		default:
			return;
		}
	}

	void Shader::Unbind( void ) const
	{
		switch( m_shaderDesc.type )
		{
		case EShaderType::VERTEX:
			m_pDevice->SetVertexShader( NULL );
			break;
		case EShaderType::PIXEL:
			m_pDevice->SetPixelShader( NULL );
			break;
		case EShaderType::EFFECT:
			break;
		default:
			return;
		}
	}

	void Shader::SetDefaults( void ) const
	{
		SIZE_T numVars = m_variables.GetSize();
		ShaderConstant *pConstant;
		
		for ( SIZE_T i = 0; i < numVars; ++i )
		{
			pConstant = m_variables[i];

			switch( pConstant->GetType() )
			{
			case EShaderConstantType::FLOAT:
				SetFloatArray( pConstant->GetName(), 
					(float*)pConstant->GetDefaultData(), 
					pConstant->GetDataCount() );
				break;
			case EShaderConstantType::INT:
				SetIntegerArray( pConstant->GetName(), 
					(int*)pConstant->GetDefaultData(), pConstant->GetDataCount() );
				break;
			case EShaderConstantType::BOOL:
				SetBooleanArray( pConstant->GetName(), 
					(bool*)pConstant->GetDefaultData(), pConstant->GetDataCount() );
				break;
			}
		}
	}

	void Shader::SetFloat( const CHAR *pName, const FLOAT *pVal )
	{
		D3DXHANDLE hConst = GetVariableByName( pName );

		if ( hConst )
			m_pConstants->SetFloat( m_pDevice, hConst, *pVal );
	}
	
	void Shader::SetFloatArray( const CHAR *pName, const float *pArray, const SIZE_T count ) const
	{
		D3DXHANDLE hConst = GetVariableByName( pName );

		if ( hConst )
			m_pConstants->SetFloatArray( m_pDevice, hConst, pArray, count );
	}

	void Shader::SetMatrix( const CHAR *pName, const Matrix4x4 *pMat )
	{
		D3DXHANDLE hConst = GetVariableByName( pName );

		if ( hConst )
			m_pConstants->SetMatrix( m_pDevice, hConst, (const D3DXMATRIX*)pMat );
	}
	
	void Shader::SetMatrixArray( const CHAR *pName, const Matrix4x4 *pArray, const SIZE_T count ) const
	{
		D3DXHANDLE hConst = GetVariableByName( pName );

		if ( hConst )
			m_pConstants->SetMatrixArray( m_pDevice, hConst, (const D3DXMATRIX*)pArray, count );
	}

	void Shader::SetInteger(const CHAR *pName, const INT *pVal )
	{
		D3DXHANDLE hConst = GetVariableByName( pName );

		if ( hConst )
			m_pConstants->SetInt( m_pDevice, hConst, *pVal );
	}
	
	void Shader::SetIntegerArray( const CHAR *pName, const INT *pArray, const SIZE_T count ) const
	{
		D3DXHANDLE hConst = GetVariableByName( pName );

		if ( hConst )
			m_pConstants->SetIntArray( m_pDevice, hConst, pArray, count );
	}

	void Shader::SetBoolean( const CHAR *pName, const BOOL *pVal )
	{
		D3DXHANDLE hConst = GetVariableByName( pName );

		if ( hConst )
			m_pConstants->SetBool( m_pDevice, hConst, *pVal );
	}
	
	void Shader::SetBooleanArray( const CHAR *pName, const BOOL *pArray, const SIZE_T count ) const
	{
		D3DXHANDLE hConst = GetVariableByName( pName );

		if ( hConst )
		{
			// NOTE: Take note that Pulse uses BOOL as int therefore we can set pArray directly. If using the c++ bool type, we can't be sure if 
			//	the size of the data type is same as int.
			m_pConstants->SetBoolArray( m_pDevice, hConst, (const ::BOOL*)pArray, count );
		}
	}

	EErrorCode::Type Shader::CreateVertexShader( void )
	{
		RendererDevice *pDevice = Renderer::GetInstance()->GetDevice();
		IDirect3DDevice9 *pNakedDevice = reinterpret_cast<IDirect3DDevice9*>(Renderer::GetInstance()->GetNakedDevice());

		ID3DXBuffer *pCode = NULL;
		ID3DXBuffer *pErrors = NULL;

		// Get Shader profile
		if ( m_shaderDesc.profile == PSX_String( "" ) )
		{
			CHAR temp[10];
			pDevice->TestCapabilities( EGraphicsCaps::VERTEX_SHADER, &temp );
			m_shaderDesc.profile = temp;
		}

		if ( m_shaderDesc.function == PSX_String( "" ) )
			m_shaderDesc.function = PSX_String( "Main" );
		
		DWORD flags = 0;

		if ( m_shaderDesc.bDebugInfo )
			flags |= D3DXSHADER_DEBUG;

		// TODO: Get macros...?

		// Load shader file
		FileHandlePtr hFile = FileSystem::GetInstance()->OpenFile( m_shaderDesc.filename.GetCString() );

		if ( !hFile->IsValid() )
		{
			PSX_PushError( "Failed to load shader file %s.", m_shaderDesc.filename );
			return EErrorCode::GRAPHICS;
		}

		ResHandlePtr hResource =  Engine::GetInstance()->GetResourceCache()->GetHandle( hFile );

		D3DXCompileShader( (LPCSTR)hResource->GetBuffer(), hResource->GetSize(), NULL, NULL, 
			m_shaderDesc.function.GetCString(), m_shaderDesc.profile.GetCString(), flags,
			&pCode, &pErrors, &m_pConstants );

		if ( pErrors || !pCode )
		{
			PSX_PushError( "Error compiling vertex shader." );
			HandleErrors( pErrors );
			PSX_SafeRelease( pErrors );
			return EErrorCode::GRAPHICS;
		}

		HRESULT hr;
		if ( FAILED( hr = pNakedDevice->CreateVertexShader( (const DWORD*)pCode->GetBufferPointer(), &m_pVertexShader )))
		{
			PSX_PushError( "Shader could not be created. %s", DX9Mapping::GetDXErrorDescription( hr ) );
			PSX_SafeRelease( pCode );
			return EErrorCode::GRAPHICS;
		}

		m_pConstants->SetDefaults( pNakedDevice );
		PSX_SafeRelease( pCode );

		return EErrorCode::NONE;
	}
	
	EErrorCode::Type Shader::CreateGeometryShader( void )
	{
		PSX_PushError( "Geometry shaders are not supported in DirectX 9.0c." );
		return EErrorCode::GRAPHICS;
	}
	
	EErrorCode::Type Shader::CreatePixelShader( void )
	{
		RendererDevice *pDevice = Renderer::GetInstance()->GetDevice();
		IDirect3DDevice9 *pNakedDevice = reinterpret_cast<IDirect3DDevice9*>(Renderer::GetInstance()->GetNakedDevice());
		
		ID3DXBuffer *pCode = NULL;
		ID3DXBuffer *pErrors = NULL;

		// Get shader profile
		if ( m_shaderDesc.profile == PSX_String("") )
		{
			CHAR temp[10];
			pDevice->TestCapabilities( EGraphicsCaps::PIXEL_SHADER, &temp );
			m_shaderDesc.profile = temp;
		}

		if ( m_shaderDesc.function == PSX_String( "" ) )
			m_shaderDesc.function = PSX_String( "Main" );

		DWORD flags = 0;

		if ( m_shaderDesc.bDebugInfo )
			flags |= D3DXSHADER_DEBUG;

		// TODO: Get macros...?

		// Load shader file
		FileHandlePtr hFile = FileSystem::GetInstance()->OpenFile( m_shaderDesc.filename.GetCString() );
		
		if ( !hFile->IsValid() )
		{
			PSX_PushError( "Failed to load shader file %s.", m_shaderDesc.filename );
			return EErrorCode::GRAPHICS;
		}

		ResHandlePtr hResource =  Engine::GetInstance()->GetResourceCache()->GetHandle( hFile );

		D3DXCompileShader( (LPCSTR)hResource->GetBuffer(), hResource->GetSize(), NULL, NULL, 
			m_shaderDesc.function.GetCString(), m_shaderDesc.profile.GetCString(), flags,
			&pCode, &pErrors, &m_pConstants );

		if ( pErrors || !pCode )
		{
			PSX_PushError( "Error compiling pixel shader." );
			HandleErrors( pErrors );
			PSX_SafeRelease( pErrors );
			return EErrorCode::GRAPHICS;
		}

		HRESULT hr;
		if ( FAILED( hr = pNakedDevice->CreatePixelShader( (const DWORD*)pCode->GetBufferPointer(), &m_pPixelShader )))
		{
			PSX_PushError( "Shader could not be created. %s", DX9Mapping::GetDXErrorDescription( hr ) );
			PSX_SafeRelease( pCode );
			return EErrorCode::GRAPHICS;
		}

		m_pConstants->SetDefaults( pNakedDevice );
		PSX_SafeRelease( pCode );

		return EErrorCode::NONE;
	}
	
	EErrorCode::Type Shader::CreateEffect( void )
	{
		RendererDevice *pDevice = Renderer::GetInstance()->GetDevice();
		IDirect3DDevice9 *pNakedDevice = reinterpret_cast<IDirect3DDevice9*>(Renderer::GetInstance()->GetNakedDevice());

		ID3DXBuffer *pErrors = NULL;

		// Load shader file
		FileHandlePtr hFile = FileSystem::GetInstance()->OpenFile( m_shaderDesc.filename.GetCString() );
		
		if ( !hFile->IsValid() )
		{
			PSX_PushError( "Failed to load shader file %s.", m_shaderDesc.filename );
			return EErrorCode::GRAPHICS;
		}

		ResHandlePtr hResource =  Engine::GetInstance()->GetResourceCache()->GetHandle( hFile );

		D3DXCreateEffect( pNakedDevice, hResource->GetBuffer(), hResource->GetSize(), NULL, NULL,
			0, NULL, &m_pEffect, &pErrors );

		if ( pErrors )
		{
			PSX_PushError( "One or more errors occured while creating a shader effect." );
			return EErrorCode::GRAPHICS;
		}

		return EErrorCode::NONE;
	}

	void Shader::HandleErrors( ID3DXBuffer *pErrors )
	{
		PSX_PushError( "Shader error: %s.", (CHAR*)pErrors->GetBufferPointer() );
	}

	D3DXHANDLE Shader::GetVariableByName( const CHAR *pName ) const
	{
		D3DXHANDLE hConst = m_pConstants->GetConstantByName( NULL, pName );
		
		if ( !hConst )
			PSX_PushError( "Variable name %s not found in shader %s", pName, m_shaderDesc.filename.GetCString() );
	
		return hConst;
	}

	void Shader::LoadConstants( void )
	{
		D3DXCONSTANTTABLE_DESC tableDesc;
		m_pConstants->GetDesc( &tableDesc );

		for ( UINT i = 0; i < tableDesc.Constants; ++i )
		{
			D3DXHANDLE constant = m_pConstants->GetConstant( NULL, i );
			D3DXCONSTANT_DESC constantDesc;
			UINT count;

			m_pConstants->GetConstantDesc( constant, &constantDesc, &count );

			// Skip textures and other objects
			if ( constantDesc.Class = D3DXPC_OBJECT )
				continue;

			// Get default values
			EShaderConstantType::Type type;
			void *pDefault = NULL;
			switch( constantDesc.Type )
			{
			case D3DXPT_FLOAT:
				if ( constantDesc.DefaultValue )
				{
					pDefault = new float[constantDesc.Rows * constantDesc.Columns * constantDesc.RegisterCount ];
					float *pInput = (float*)pDefault;
					float *pBuf = (float*)constantDesc.DefaultValue;
					for ( UINT i = 0; i < constantDesc.RegisterCount; ++i )
					{
						for ( UINT j = 0; j < constantDesc.Rows; ++j )
						{
							PSX_MemCopy( pInput, pBuf, constantDesc.Columns * sizeof(float) );
							pInput += constantDesc.Columns;
							pBuf += sizeof(float);
						}
					}
				}
				type = EShaderConstantType::FLOAT;
				break;
			case D3DXPT_INT:
				if ( constantDesc.DefaultValue )
				{
					pDefault = new int[constantDesc.Rows * constantDesc.Columns * constantDesc.RegisterCount];
					int* pInput = (int*)pDefault;
					int* pBuf = (int*)constantDesc.DefaultValue;
					for( UINT j = 0; j < constantDesc.RegisterCount; ++j )
					{
						for( UINT i = 0; i < constantDesc.Rows; ++i )
						{
							PSX_MemCopy( pInput, pBuf, constantDesc.Columns * sizeof(int) );
							pInput += constantDesc.Columns;
							pBuf += sizeof(int);
						}
					}
				}

				type = EShaderConstantType::INT;
				break;
			case D3DXPT_BOOL:
				if(constantDesc.DefaultValue)
				{
					pDefault = new bool[constantDesc.Rows * constantDesc.Columns * constantDesc.RegisterCount];
					bool* pInput = (bool*)pDefault;
					int* pBuf = (int*)constantDesc.DefaultValue;
					for( UINT j = 0; j < constantDesc.RegisterCount; ++j )
					{
						for( UINT i = 0; i < constantDesc.Rows; ++i )
						{
							PSX_MemSet( pInput, (*pBuf ? true : false), constantDesc.Columns * sizeof(bool) );
							pInput += constantDesc.Columns;
							pBuf += sizeof(int);
						}
					}
				}
				type = EShaderConstantType::BOOL;
				break;

			default:
				if(constantDesc.DefaultValue)
				{
					pDefault = new float[constantDesc.Rows * constantDesc.Columns * constantDesc.RegisterCount];
					float* pInput = (float*)pDefault;
					float* pBuf = (float*)constantDesc.DefaultValue;
					for( UINT j = 0; j < constantDesc.RegisterCount; ++j )
					{
						for( UINT i = 0; i < constantDesc.Rows; ++i )
						{
							PSX_MemCopy( pInput, pBuf, constantDesc.Columns * sizeof(float) );
							pInput += constantDesc.Columns;
							pBuf += sizeof(float);
						}
					}
				}
				type = EShaderConstantType::FLOAT;
				break;
			}
		
			SIZE_T constantDataCount = constantDesc.Elements * constantDesc.Columns * constantDesc.Rows;
			ShaderConstant *pShaderConst = new ShaderConstant( constantDesc.Name, type, constantDataCount, pDefault );

			m_variables.PushBack( pShaderConst );
			m_varMap[ constantDesc.Name ] = pShaderConst;

			PSX_SafeDeleteArray( pDefault );	
		}
	}

}

#endif __PSX_WINDOWS_WINDOW
